package linker

import (
	"debug/elf"
	"risc-v/ld/pkg/utils"
	"sort"
)

// MergedSection 合并后的节
// 把多个 MergeableSection 合并为一个 MergedSection
type MergedSection struct {
	Chunk
	Map map[string]*SectionFragment
}

func NewMergedSection(name string, flags uint64, typ uint32) *MergedSection {
	m := &MergedSection{
		Chunk: NewChunk(),
		Map:   make(map[string]*SectionFragment),
	}

	m.Name = name
	m.SectionHeader.Flags = flags
	m.SectionHeader.Type = typ
	return m
}

// Insert 向 MergedSection 中插入一个 SectionFragment
func (m *MergedSection) Insert(key string, p2align uint32) *SectionFragment {
	fragment, ok := m.Map[key]
	if !ok {
		fragment = NewSectionFragment(m)
		m.Map[key] = fragment
	}

	// 取最大对齐量
	if fragment.P2Align < p2align {
		fragment.P2Align = p2align
	}

	return fragment
}

// GetMergedSectionInstance 获取一个合并节实例
// 存在则返回，不存在则创建
// 根据传进来的 InputSection 的信息来匹配对应的 MergedSection
func GetMergedSectionInstance(ctx *Context, name string, typ uint32, flags uint64) *MergedSection {
	// 获取对应的 输出节（合并节）的名字
	name = GetOutPutName(name, flags)
	// 清除对应的位
	flags = flags &^ uint64(elf.SHF_GROUP) &^ uint64(elf.SHF_MERGE) &^ uint64(elf.SHF_STRINGS) &^ uint64(elf.SHF_COMPRESSED)

	find := func() *MergedSection {
		for _, outSec := range ctx.MergedSections {
			// 该InputSection 存在这个 MergedSection 中
			if outSec.Name == name && outSec.SectionHeader.Flags == flags && outSec.SectionHeader.Type == typ {
				return outSec
			}
		}
		return nil
	}

	if outSec := find(); outSec != nil {
		return outSec
	}
	outSec := NewMergedSection(name, flags, typ)
	ctx.MergedSections = append(ctx.MergedSections, outSec)

	return outSec
}

// AssignOffsets 计算合并节的偏移
// 遍历所有 sectionFragment 计算单个segmentFragment 大小， 对齐之后不断累加
func (m *MergedSection) AssignOffsets() {

	var fragments []struct {
		Key string
		Val *SectionFragment
	}

	for key := range m.Map {
		fragments = append(fragments, struct {
			Key string
			Val *SectionFragment
		}{
			Key: key,
			Val: m.Map[key],
		})
	}

	// 对齐量小的排序在前,可以减小可执行文件的体积
	sort.SliceStable(fragments, func(i, j int) bool {
		x := fragments[i]
		y := fragments[j]

		if x.Val.P2Align != y.Val.P2Align {
			return x.Val.P2Align < y.Val.P2Align
		}

		if len(x.Key) != len(y.Key) {
			return len(x.Key) < len(y.Key)
		}

		return x.Key < y.Key
	})

	offset := uint64(0)
	p2align := uint64(0)

	for _, frag := range fragments {
		offset = utils.AlignTo(offset, 1<<frag.Val.P2Align)
		frag.Val.Offset = uint32(offset)
		offset += uint64(len(frag.Key))

		if p2align < uint64(frag.Val.P2Align) {
			p2align = uint64(frag.Val.P2Align)
		}
	}

	m.SectionHeader.Size = utils.AlignTo(offset, 1<<p2align)
	m.SectionHeader.AddrAlign = 1 << p2align
}

func (m *MergedSection) CopyBuf(ctx *Context) {
	buf := ctx.Buf[m.SectionHeader.Offset:]
	for key := range m.Map {
		if frag, ok := m.Map[key]; ok {
			copy(buf[frag.Offset:], key)
		}
	}
}
